import copy
import random
from Tools.Paraphraser import Paraphraser
from Tools.PlaceholderUtils import extract_placeholders
from Tools.Translator import Translator
from Logger import logger


def random_delete(genetic_prompt):
    """随机删除一个chunk"""
    if len(genetic_prompt.sequence) == 0:
        return
    for _ in range(3):
        index = random.choice(genetic_prompt.sequence)
        placeholders = extract_placeholders(genetic_prompt.chunks[index])
        if placeholders:
            new_chunk = '\n'.join(placeholders)
            if new_chunk == genetic_prompt.chunks[index]:
                continue
            else:
                logger.debug(f"Delete: {genetic_prompt.chunks[index]} -> {new_chunk}")
                genetic_prompt.chunks[index] = new_chunk
                break
        logger.debug(f"Delete: {genetic_prompt.chunks[index]}")
        genetic_prompt.chunks[index] = ""
        break


class random_paraphrase:
    def __init__(self, paraphraser):
        self.paraphraser = paraphraser

    def __call__(self, genetic_prompt):
        """随机替换一个chunk"""
        if len(genetic_prompt.sequence) == 0:
            return
        index = random.choice(genetic_prompt.sequence)
        new_chunk = self.paraphraser.run(genetic_prompt.chunks[index])
        logger.debug(f"Paraphrase: {genetic_prompt.chunks[index]} -> {new_chunk}")
        genetic_prompt.chunks[index] = new_chunk


class random_translate:
    def __init__(self, translator):
        self.translator = translator

    def __call__(self, genetic_prompt):
        """随机替换一个chunk"""
        if len(genetic_prompt.sequence) == 0:
            return
        index = random.choice(genetic_prompt.sequence)
        new_chunk = self.translator.run(genetic_prompt.chunks[index])
        logger.debug(f"Translate: {genetic_prompt.chunks[index]} -> {new_chunk}")
        genetic_prompt.chunks[index] = new_chunk


def random_exchange(genetic_prompt):
    """随机交换两个chunk"""
    if len(genetic_prompt.sequence) < 2:
        return
    index1, index2 = random.sample(genetic_prompt.sequence, 2)
    logger.debug(f"Exchange: {genetic_prompt.chunks[index1]} <-> {genetic_prompt.chunks[index2]}")
    genetic_prompt.sequence[index1], genetic_prompt.sequence[index2] = genetic_prompt.sequence[index2], \
    genetic_prompt.sequence[index1]


class random_expand:
    def __init__(self, paraphraser, translator):
        self.paraphraser = paraphraser
        self.translator = translator

    def __call__(self, genetic_prompt):
        """重复某个chunk"""
        if len(genetic_prompt.sequence) == 0:
            return
        gene = random.choice(genetic_prompt.chunks)
        fn = random.choice([self.paraphraser.run, self.translator.run])
        new_sentence = fn(gene)
        genetic_prompt.sequence.append(len(genetic_prompt.chunks))
        genetic_prompt.chunks.append(new_sentence)


class Mutator:
    def __init__(
            self,
            args,
    ):
        paraphraser = Paraphraser(args.paraphraser_prompt, args.paraphraser_model)
        translator = Translator(args.translator_prompt, args.translator_model)

        self.props = [
            args.deletion_prob,
            args.paraphrasing_prob,
            args.translation_prob,
            args.expansion_prob,
            args.exchanging_prob
        ]
        self.mutation_list = [
            random_delete,
            random_paraphrase(paraphraser),
            random_translate(translator),
            random_expand(paraphraser, translator),
            random_exchange
        ]

    def run(self, genetic_prompt):
        mutation = random.choices(self.mutation_list, weights=self.props, k=1)[0]
        new_genetic_prompt = copy.deepcopy(genetic_prompt)
        mutation(new_genetic_prompt)
        new_genetic_prompt.clear_scores()
        return new_genetic_prompt
